#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time: 2017/08/18 上午11:21
@Author: CZC
@File: connectdb.py
"""

import pymongo
from common.logger import global_log


class MongoDB(object):
    def __init__(self, mongodb):
        self.mongodb = mongodb
        self.client = pymongo.MongoClient(self.mongodb.MONGODB_SERVER, self.mongodb.MONGODB_PORT)
        self.db = self.client[self.mongodb.MONGODB_DB]
        self.db.authenticate(self.mongodb.MONGODB_USER, self.mongodb.MONGODB_PSW)
        self.log = global_log

        # 建立索引
        # self.collection.ensure_index([("_utime", pymongo.ASCENDING)])
        # self.collection.ensure_index([("_in_time", pymongo.ASCENDING)])
        # self.collection.ensure_index([("_record_id", pymongo.ASCENDING)])

    def __del__(self):
        self.log.info('mongodb client close……')
        self.client.close()

    def traverse_batch(self, table, where=None):
        cursor = None
        try:
            where = {} if where is None else where
            cursor = self.db[table].find(where, no_cursor_timeout=True).batch_size(500)
            for item in cursor:
                yield item
        except Exception as e:
            self.log.exception(e)
            raise e
        finally:
            if cursor is not None:
                cursor.close()

    def find_data(self, table, where=None):
        cursor = None
        try:
            where = {} if where is None else where
            cursor = self.db[table].find(where, no_cursor_timeout=True)
            for item in cursor:
                yield item
        except Exception as e:
            self.log.exception(e)
            raise e
        finally:
            if cursor is not None:
                cursor.close()

    def find_and_modify(self, table, query=None, update=None,
                        upsert=False, sort=None, full_response=False,
                        manipulate=False, **kwargs):
        query = {} if query is None else query
        try:
            self.db[table].find_and_modify(query=query, update=update, upsert=upsert,
                                           sort=sort, full_response=full_response,
                                           manipulate=manipulate, **kwargs)
        except Exception as e:
            self.log.exception(e)
            raise e

    # def find_data(self, value, collection):
    #     coll = self.db[collection]
    #     cursor = coll.find({'_id': value}, no_cursor_timeout=True)
    #     data_list = list(cursor)
    #     if data_list:
    #         return True
    #     return False

    def get_keyword(self, province, collection):
        coll = self.db[collection]
        cursor = coll.find({}, no_cursor_timeout=True)
        for doc in cursor:
            keyword = '{}{}'.format(province, doc.get('_id'))
            yield keyword

    def get_company(self, collection):
        coll = self.db[collection]
        cursor = coll.find({}, no_cursor_timeout=True)
        for doc in cursor:
            keyword = '{}'.format(doc.get('_id'))
            yield keyword

    def insert_items(self, value, collection):
        coll = self.db[collection]
        # coll.insert({'_id': value})

        cursor = coll.find({'_id': value})
        cursor = list(cursor)
        if not cursor:
            coll.insert({'_id': value})

